#!/bin/bash
set -e

# 验证 torch 是否可用（xformers 构建时需要）
echo "验证 PyTorch 是否可用..."
python3 -c "import torch; print(f'PyTorch 版本: {torch.__version__}'); print(f'PyTorch 路径: {torch.__file__}'); print(f'使用设备: CPU')" || {
    echo "错误: PyTorch 不可用，尝试安装..."
    # 如果 torch 不可用，先安装 CPU 版本
    # 先尝试从 PyTorch 官方源安装
    python3 -m pip install --no-cache-dir \
        torch torchvision torchaudio \
        --index-url https://download.pytorch.org/whl/cpu || \
    python3 -m pip install --no-cache-dir \
        torch torchvision torchaudio \
        -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn || {
        echo "错误: 无法安装 PyTorch，请检查环境和网络连接"
        exit 1
    }
    # 再次验证
    python3 -c "import torch; print(f'PyTorch 安装成功，版本: {torch.__version__}')" || {
        echo "错误: PyTorch 安装后仍无法导入"
        exit 1
    }
}

# 如果设置了 LOCAL_MODEL_PATH 且存在 requirements.txt，安装依赖
if [ -n "$LOCAL_MODEL_PATH" ] && [ -f "$LOCAL_MODEL_PATH/requirements.txt" ]; then
    echo "发现 AudioFly 项目 requirements.txt，安装依赖..."
    
    # 直接安装所有依赖，包括 xformers
    python3 -m pip install --no-cache-dir \
        -r "$LOCAL_MODEL_PATH/requirements.txt" \
        -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn || {
        echo "警告: 部分依赖安装失败，尝试安装 xformers 的构建依赖后重试..."
        # 如果安装失败，可能是 xformers 需要构建依赖
            python3 -m pip install --no-cache-dir \
                ninja \
                packaging \
                wheel \
                -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
            
        # 再次尝试安装所有依赖
            python3 -m pip install --no-cache-dir \
            -r "$LOCAL_MODEL_PATH/requirements.txt" \
                -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn || {
            echo "错误: 依赖安装失败，请检查 requirements.txt 和构建环境"
                exit 1
        }
    }
    
    # 验证 xformers 是否安装成功
    python3 -c "import xformers; print(f'xformers 版本: {xformers.__version__}')" || {
        echo "警告: xformers 未安装或无法导入，AudioFly 可能需要 xformers"
    }
fi

# 启动应用
exec python3 -m uvicorn api.main:app --host 0.0.0.0 --port 8000 --app-dir /app

